#include <iostream>

extern "C" void hello_plugin() {
    std::cout << "Hello from YOLO Plugin!" << std::endl;
}

// 这是 YOLO 层的基本实现
__global__ void yolo_kernel(float* input, float* output, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        output[idx] = input[idx] * 2.0f;  // 示例操作
    }
}

extern "C" void launch_yolo_kernel(float* input, float* output, int size) {
    dim3 block(256);
    dim3 grid((size + block.x - 1) / block.x);
    yolo_kernel<<<grid, block>>>(input, output, size);
    cudaDeviceSynchronize();
}
